"""
评估模块
"""

import torch
from torch.utils.data import DataLoader
from dataset import get_dataloaders
from model import get_model
from utils import load_checkpoint, accuracy
from config import Config
import matplotlib.pyplot as plt
import numpy as np
import os
from .dataset import *

def evaluate_model():
    """
    评估模型性能
    """
    cfg = Config()
    
    # 获取数据和模型
    _, test_loader = get_dataloaders()
    model = get_model()
    
    # 加载最佳模型
    load_checkpoint(model)
    model.eval()
    
    # 评估
    test_loss = 0
    correct = 0
    total = 0
    criterion = torch.nn.CrossEntropyLoss()
    
    # 用于计算每个类的准确率
    class_correct = list(0. for _ in range(cfg.NUM_CLASSES))
    class_total = list(0. for _ in range(cfg.NUM_CLASSES))
    
    with torch.no_grad():
        for inputs, targets in test_loader:
            inputs, targets = inputs.to(cfg.DEVICE), targets.to(cfg.DEVICE)
            outputs = model(inputs)
            loss = criterion(outputs, targets)
            
            test_loss += loss.item()
            _, predicted = outputs.max(1)
            total += targets.size(0)
            correct += predicted.eq(targets).sum().item()
            
            # 计算每个类的准确率
            c = (predicted == targets).squeeze()
            for i in range(targets.size(0)):
                label = targets[i]
                class_correct[label] += c[i].item()
                class_total[label] += 1
    
    # 打印总体结果
    test_loss /= len(test_loader)
    test_acc = 100. * correct / total
    print(f'Test Loss: {test_loss:.4f}, Test Acc: {test_acc:.2f}%')
    
    # 打印每个类的准确率
    superclass_mapping = get_cifar100_superclass_mapping()
    for i in range(cfg.NUM_CLASSES):
        class_name = ', '.join(superclass_mapping[i])[:30] + '...'  # 截断长名称
        print(f'Accuracy of {class_name:30s} : {100 * class_correct[i] / (class_total[i] + 1e-5):.2f}%')
    
    # 可视化混淆矩阵
    plot_confusion_matrix(model, test_loader, cfg)

def plot_confusion_matrix(model, test_loader, cfg):
    """
    绘制混淆矩阵
    """
    from sklearn.metrics import confusion_matrix
    import seaborn as sns
    
    model.eval()
    all_preds = []
    all_targets = []
    
    with torch.no_grad():
        for inputs, targets in test_loader:
            inputs = inputs.to(cfg.DEVICE)
            outputs = model(inputs)
            _, preds = torch.max(outputs, 1)
            
            all_preds.extend(preds.cpu().numpy())
            all_targets.extend(targets.cpu().numpy())
    
    # 计算混淆矩阵
    cm = confusion_matrix(all_targets, all_preds)
    
    # 绘制
    plt.figure(figsize=(12, 10))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues')
    plt.xlabel('Predicted')
    plt.ylabel('True')
    plt.title('Confusion Matrix')
    plt.savefig(os.path.join(cfg.LOG_DIR, 'confusion_matrix.png'))
    plt.close()